#include "naive_neighbor.hpp"

#include "MG.hpp"
#include "parcsr_matop.hpp"
#include <queue>

template<typename idx_t, typename setup_t, int dof>
idx_t* naive_neighbor(const par_CSRMatrix<idx_t, setup_t, setup_t, dof> * par_fine_mat,
    par_CSRMatrix<idx_t, setup_t, setup_t, 1> * & R_mat, par_CSRMatrix<idx_t, setup_t, setup_t, 1> * & P_mat)
{
    const MPI_Comm comm = par_fine_mat->comm;
    int my_pid; MPI_Comm_rank(comm, & my_pid);
    int num_procs; MPI_Comm_size(comm, & num_procs);
    constexpr MPI_Datatype mpi_idx_type = sizeof(idx_t) == 8 ? MPI_LONG_LONG : MPI_INT;

    const CSRMatrix<idx_t, setup_t, setup_t, dof> & diag = par_fine_mat->diag;
    const idx_t my_num_F = diag.nrows; assert(my_num_F == par_fine_mat->end_row - par_fine_mat->beg_row);
    idx_t my_num_C = 0;

    idx_t * loc_map_f2c = new idx_t [my_num_F];
    std::queue<idx_t> to_examin;
    bool * been_in_Q = new bool [my_num_F];// 记录一个点是否已经入过队检查
    coar_wells.clear();// 为这一轮粗化清空
    constexpr idx_t UNAGGREGATED = -1;
    // 初始化
    #pragma omp parallel for schedule(static)
    for (idx_t fi = 0; fi < my_num_F; fi++) {
        loc_map_f2c[fi] = UNAGGREGATED;
        been_in_Q  [fi] = false;
    }
    
    idx_t last_found = -1;
    auto find_next_unexamined = [& diag, & been_in_Q, my_num_F, & last_found] () {
        idx_t i = last_found + 1;
        for ( ; i < my_num_F; i ++) {
            if (been_in_Q[i] == false && // 之前没进过队
                diag.row_ptr[i + 1] - diag.row_ptr[i] > 1)// 有不止一个（即自己）的邻居
            {
                been_in_Q[i] = true;
                break;
            }
        }
        last_found = i;
        return i;
    };

    while (true) {
        if (to_examin.empty()) {
            const idx_t next2exam = find_next_unexamined();
            if (next2exam == my_num_F) break;
            else to_examin.push(next2exam);
        }

        idx_t loc_fi = to_examin.front(); assert(been_in_Q[loc_fi] == true);
        to_examin.pop();
        // if (my_pid == 273) printf("check loc_fi %d :: %d num_loc_ngbs %d\n", loc_fi, loc_map_f2c[loc_fi], diag.row_ptr[loc_fi + 1] - diag.row_ptr[loc_fi]);

        for (idx_t p = diag.row_ptr[loc_fi]; p < diag.row_ptr[loc_fi + 1]; p ++) {// 将所有未检查过的邻居加入待检查队列
            const idx_t loc_fj = diag.col_idx[p];
            if (been_in_Q[loc_fj] == false) {
                to_examin.push(loc_fj);// 加入待检查的队列
                been_in_Q[loc_fj] = true;
            }
        }

        if (loc_map_f2c[loc_fi] != UNAGGREGATED) {// 之前已经聚合过了
            continue;
        }
        if (iso_wells.find(loc_fi + par_fine_mat->beg_row) != iso_wells.end()) {// 井不与别人聚合，单独成一个粗点
            assert(loc_map_f2c[loc_fi] == UNAGGREGATED);// 之前不能被聚合过
            coar_wells.emplace(my_num_C);
            loc_map_f2c[loc_fi] = my_num_C;
            my_num_C ++;
            continue;
        }

        bool can_be_root = true;// 是否可以作为聚合的root point，要求所有邻居都未被聚合
        for (idx_t p = diag.row_ptr[loc_fi]; p < diag.row_ptr[loc_fi + 1] && can_be_root; p ++) {
            const idx_t loc_fj = diag.col_idx[p];
            can_be_root =   loc_map_f2c[loc_fj] == UNAGGREGATED || // 该邻居应该未聚合
                            iso_wells.find(loc_fj + par_fine_mat->beg_row) != iso_wells.end();// 如果已聚合，那只能是孤立点
        }
        // if (my_pid == 273) printf("Proc %d loc_fi %d num_ngbs %d root %d\n", my_pid, loc_fi, diag.row_ptr[loc_fi + 1] - diag.row_ptr[loc_fi], can_be_root);
        if (can_be_root == false) continue;// 该点不能作为root point聚合它的邻居
        // 否则，可以记录聚合信息
        for (idx_t p = diag.row_ptr[loc_fi]; p < diag.row_ptr[loc_fi + 1]; p ++) {
            const idx_t loc_fj = diag.col_idx[p];
            if (iso_wells.find(loc_fj + par_fine_mat->beg_row) == iso_wells.end()) {
                loc_map_f2c[loc_fj] = my_num_C;
                // printf("Map F %d to C %d\n", loc_fj, my_num_C);
            }
        }
        my_num_C ++;
    }
    assert(last_found == my_num_F);

    // for (int p = 0; p < num_procs; p ++) {
    //     if ( p == my_pid ) {
    //         printf("Proc %d num_F %d num_C %d\n", my_pid, my_num_F, my_num_C);
    //     }
    //     MPI_Barrier(comm);
    // }

    constexpr int e_size = dof*dof;
    auto calc_blk_norm = [] (const setup_t * val) {
        setup_t ret = 0.0;
        for (int f = 0; f < e_size; f++) ret += val[f] * val[f];
        return sqrt(ret);
    };

    idx_t * C_rows_partits = new idx_t [num_procs + 1];
    C_rows_partits[0] = 0;
    MPI_Allgather(& my_num_C      , 1, sizeof(idx_t) == 8 ? MPI_LONG_LONG : MPI_INT,
                C_rows_partits + 1, 1, sizeof(idx_t) == 8 ? MPI_LONG_LONG : MPI_INT, comm);
    for (int p = 0; p < num_procs; p++)
        C_rows_partits[p + 1] += C_rows_partits[p];
    
    // 构造限制矩阵
    R_mat = new par_CSRMatrix<idx_t, setup_t, setup_t, 1>(comm, C_rows_partits, par_fine_mat->rows_partition);
    P_mat = new par_CSRMatrix<idx_t, setup_t, setup_t, 1>(comm, par_fine_mat->cols_partition, C_rows_partits);
#if 1 // 允许跨进程聚合
    const idx_t my_F_beg = par_fine_mat->beg_row;
    const idx_t _cbeg = C_rows_partits[my_pid], _cend = C_rows_partits[my_pid + 1];
    idx_t * glb_map_f2c = new idx_t [my_num_F];// loc_f => glb_c
    for (idx_t loc_fi = 0; loc_fi < my_num_F; loc_fi ++) {
        if (loc_map_f2c[loc_fi] >= 0)
            glb_map_f2c[loc_fi] = loc_map_f2c[loc_fi] + _cbeg;// 累加全局偏移
        else
            glb_map_f2c[loc_fi] = -1;// 暂时置为-1
    }
    {// 处理尚未聚合的点
        idx_t * send_buf = new idx_t [par_fine_mat->commpkg->send_map_starts[par_fine_mat->commpkg->num_sends]];
        idx_t * recv_buf = new idx_t [par_fine_mat->offd.ncols];
        for (idx_t s = 0; s < par_fine_mat->commpkg->num_sends; s++) {
            idx_t       * buf    = send_buf                              + par_fine_mat->commpkg->send_map_starts[s];
            const idx_t * my_ids = par_fine_mat->commpkg->send_map_elmts + par_fine_mat->commpkg->send_map_starts[s];
            const idx_t num = par_fine_mat->commpkg->send_map_starts[s + 1] - par_fine_mat->commpkg->send_map_starts[s];
            for (idx_t k = 0; k < num; k++) {
                buf[k] = glb_map_f2c[my_ids[k]];
            }
            MPI_Isend(buf, num, mpi_idx_type, par_fine_mat->commpkg->send_pids[s], 114, comm, par_fine_mat->commpkg->send_reqs + s);
        }
        for (idx_t r = 0; r < par_fine_mat->commpkg->num_recvs; r++) {
            idx_t * buf = recv_buf + par_fine_mat->commpkg->recv_vec_starts[r];
            const idx_t num = par_fine_mat->commpkg->recv_vec_starts[r + 1] - par_fine_mat->commpkg->recv_vec_starts[r];
            MPI_Irecv(buf, num, mpi_idx_type, par_fine_mat->commpkg->recv_pids[r], 114, comm, par_fine_mat->commpkg->recv_reqs + r);
        }
        MPI_Waitall(par_fine_mat->commpkg->num_recvs, par_fine_mat->commpkg->recv_reqs, MPI_STATUSES_IGNORE);
    
        for (idx_t loc_fi = 0; loc_fi < my_num_F; loc_fi++) {
            if (loc_map_f2c[loc_fi] == UNAGGREGATED) {// 该点仍未被聚合
                // 从所有邻居中选最大的进行聚合：注意包含落在别的进程上的邻居
                // 找出diag内最强耦合的邻居
                const idx_t num_diag =  par_fine_mat->diag.row_ptr[loc_fi + 1] - 
                                        par_fine_mat->diag.row_ptr[loc_fi    ]  ;
                setup_t diag_vbuf [ num_diag ];
                idx_t   diag_ibuf [ num_diag ];
                for (idx_t p = par_fine_mat->diag.row_ptr[loc_fi]; p < par_fine_mat->diag.row_ptr[loc_fi + 1]; p ++) {
                    const idx_t glb_fj = par_fine_mat->diag.col_idx[p] + my_F_beg;
                    if (iso_wells.find(glb_fj) == iso_wells.end()) {// 正常点计算范数
                        diag_vbuf[p - par_fine_mat->diag.row_ptr[loc_fi]] = calc_blk_norm(par_fine_mat->diag.vals + p * e_size);
                    } else {// 否则该细点为独立点，为避免与它聚合，联通值设为0
                        diag_vbuf[p - par_fine_mat->diag.row_ptr[loc_fi]] = 0.0;
                    }
                    diag_ibuf[p - par_fine_mat->diag.row_ptr[loc_fi]] = glb_fj;
                }
#ifdef PRES_COL_ORDER // 保证列号升序（以便coupled和naive在单进程时严格复现）
                sort_vals_with_indices<idx_t, setup_t>(num_diag, diag_vbuf, diag_ibuf);
#else // 快排不稳定排序
                my_qsort1<idx_t, setup_t, idx_t, -1>(diag_vbuf, diag_ibuf, 0, num_diag - 1);
#endif
                idx_t diag_candi = -1; setup_t diag_v = 0.0;// idx_t dc_ptr = -1;
                for (idx_t p = 0; p < num_diag; p ++) {// 从大到小考察邻居
                    const idx_t glb_fj = diag_ibuf[p];
                    const idx_t loc_fj = glb_fj - my_F_beg;
                    if (loc_fj != loc_fi && // 该邻居不能是自己
                        iso_wells.find(glb_fj) == iso_wells.end() && // 也不能是井
                        glb_map_f2c[loc_fj] != UNAGGREGATED) // 其已经聚合
                    {
                        diag_candi = glb_map_f2c[loc_fj];// 该邻居所属的
                        diag_v     = diag_vbuf[p];
                        break;
                    }
                }
                // 找出offd内最强耦合的邻居
                const idx_t num_offd =  par_fine_mat->offd.row_ptr[loc_fi + 1] -
                                        par_fine_mat->offd.row_ptr[loc_fi    ]  ;
                setup_t offd_vbuf[ num_offd ];
                idx_t   offd_ibuf[ num_offd ];
                for (idx_t p = par_fine_mat->offd.row_ptr[loc_fi]; p < par_fine_mat->offd.row_ptr[loc_fi + 1]; p ++) {
                    const idx_t glb_fj = par_fine_mat->col_map_offd[ par_fine_mat->offd.col_idx[p] ];
                    if (iso_wells.find(glb_fj) == iso_wells.end()) {// 正常点计算范数
                        offd_vbuf[p - par_fine_mat->offd.row_ptr[loc_fi]] = calc_blk_norm(par_fine_mat->offd.vals + p * e_size);
                    } else {// 否则该细点为独立点，为避免与它聚合，联通值设为0
                        offd_vbuf[p - par_fine_mat->offd.row_ptr[loc_fi]] = 0.0;
                    }
                    offd_ibuf[p - par_fine_mat->offd.row_ptr[loc_fi]] = par_fine_mat->offd.col_idx[p];
                }
#ifdef PRES_COL_ORDER // 保证列号升序（以便coupled和naive在单进程时严格复现）
                sort_vals_with_indices<idx_t, setup_t>(num_offd, offd_vbuf, offd_ibuf);
#else // 快排不稳定排序
                my_qsort1<idx_t, setup_t, idx_t, -1>(offd_vbuf, offd_ibuf, 0, num_offd - 1);
#endif
                idx_t offd_candi = -1; setup_t offd_v = 0.0;// idx_t oc_ptr = -1;
                for (idx_t p = 0; p < num_offd; p ++) {// 从大到小考察邻居
                    const idx_t loc_fj = offd_ibuf[p];
                    const idx_t glb_cj = recv_buf[loc_fj];
                    if (glb_cj != -1)// 不能是邻居进程里尚未聚合的点
                    {
                        offd_candi = glb_cj;
                        offd_v     = offd_vbuf[p];
                        break;
                    }
                }
                idx_t final_candi = -1;// 最终确定的聚合邻居
                if (diag_candi != -1 && offd_candi != -1) {
                    if (diag_v > offd_v) final_candi = diag_candi;// 选值大的
                    else                 final_candi = offd_candi;
                }
                else { final_candi = MAX(diag_candi, offd_candi); }
                glb_map_f2c[loc_fi] = final_candi;
                if (final_candi == -1) {// check: 此时还尚未被聚合的，一定是只有一个邻居（即自己的）
                    assert(par_fine_mat->diag.row_ptr[loc_fi + 1] - par_fine_mat->diag.row_ptr[loc_fi] == 1);
                    assert(par_fine_mat->diag.col_idx[par_fine_mat->diag.row_ptr[loc_fi]] == loc_fi);
                }
            }
        }
        MPI_Waitall(par_fine_mat->commpkg->num_sends, par_fine_mat->commpkg->send_reqs, MPI_STATUSES_IGNORE);
        delete [] recv_buf; delete [] send_buf;
    }// 此时在glb_map_f2c[]中仍为-1的点将不被粗化进入下一层粗网格

    {// 构造插值矩阵
        P_mat->diag.nrows = my_num_F; P_mat->offd.nrows = my_num_F;
        P_mat->diag.row_ptr = new idx_t [my_num_F + 1]; P_mat->diag.row_ptr[0] = 0;
        P_mat->offd.row_ptr = new idx_t [my_num_F + 1]; P_mat->offd.row_ptr[0] = 0;
        std::unordered_map<idx_t, std::pair<int, idx_t> > Poffd_g2l;// 记录插值矩阵的非对角部分的全局序到局部序的映射
        for (idx_t loc_fi = 0; loc_fi < my_num_F; loc_fi ++) {
            const idx_t glb_ci = glb_map_f2c[loc_fi];
            if (_cbeg <= glb_ci && glb_ci < _cend) {// 粗点属于本进程负责计算
                P_mat->diag.row_ptr[loc_fi + 1] = 1;
                P_mat->offd.row_ptr[loc_fi + 1] = 0;
            } else if (glb_ci >= 0) {// 粗点不属于本进程计算
                P_mat->diag.row_ptr[loc_fi + 1] = 0;
                P_mat->offd.row_ptr[loc_fi + 1] = 1;
                
                std::pair<int, idx_t> tmp(-1, -1);
                Poffd_g2l.emplace(glb_ci, tmp);
            } else { assert(glb_ci == -1);// 粗点不进入下一层
                P_mat->diag.row_ptr[loc_fi + 1] = 0;
                P_mat->offd.row_ptr[loc_fi + 1] = 0;
            }
        }
        for (idx_t loc_fi = 0; loc_fi < my_num_F; loc_fi ++) {// prefix sum
            P_mat->diag.row_ptr[loc_fi + 1] += P_mat->diag.row_ptr[loc_fi];
            P_mat->offd.row_ptr[loc_fi + 1] += P_mat->offd.row_ptr[loc_fi];
        }
        P_mat->prepare_commpkg(Poffd_g2l);

        // 确定列数和非零元数
        P_mat->diag.ncols = my_num_C;                      P_mat->offd.ncols = Poffd_g2l.size();
        P_mat->diag.nnz   = P_mat->diag.row_ptr[my_num_F]; P_mat->offd.nnz   = P_mat->offd.row_ptr[my_num_F];
        // 分配列序号和值的内存
        P_mat->diag.col_idx = new idx_t  [P_mat->diag.nnz]; P_mat->offd.col_idx = new idx_t  [P_mat->offd.nnz];
        P_mat->diag.vals    = new setup_t[P_mat->diag.nnz]; P_mat->offd.vals    = new setup_t[P_mat->offd.nnz];

        for (idx_t loc_fi = 0; loc_fi < my_num_F; loc_fi ++) {
            const idx_t glb_ci = glb_map_f2c[loc_fi];
            // const idx_t dcnt = 0, ocnt = 0;
            const idx_t doff = P_mat->diag.row_ptr[loc_fi],
                        ooff = P_mat->offd.row_ptr[loc_fi];
            if (_cbeg <= glb_ci && glb_ci < _cend) {// 粗点属于本进程负责计算
                P_mat->diag.col_idx[doff] = glb_ci - _cbeg;
                P_mat->diag.vals   [doff] = 1.0;
            } else if (glb_ci >= 0) {// 粗点不属于本进程计算
                assert(Poffd_g2l.find(glb_ci) != Poffd_g2l.end());
                P_mat->offd.col_idx[ooff] = Poffd_g2l[glb_ci].second;
                P_mat->offd.vals   [ooff] = 1.0;
            } else { assert(glb_ci == -1);// 粗点不进入下一层
                // nothing to do
            }
        }
#ifdef DEBUG
        for (int p = 0; p < num_procs; p++) {
            if (my_pid == p) {
                printf("Proc %d Prlg diag %d %d %d offd %d %d %d\n", my_pid, P_mat->diag.nrows, P_mat->diag.ncols, P_mat->diag.nnz,
                    P_mat->offd.nrows, P_mat->offd.ncols, P_mat->offd.nnz);
                for (idx_t i = 0; i < P_mat->diag.nrows; i++) {
                    printf("  %d : ", P_mat->beg_row + i);
                    for (idx_t p = P_mat->diag.row_ptr[i]; p < P_mat->diag.row_ptr[i + 1]; p++)
                        printf("%d ", P_mat->beg_col + P_mat->diag.col_idx[p]);
                    printf(" | ");
                    for (idx_t p = P_mat->offd.row_ptr[i]; p < P_mat->offd.row_ptr[i + 1]; p++)
                        printf("%d ", P_mat->col_map_offd[P_mat->offd.col_idx[p]]);
                    printf("\n");
                }
            }
            MPI_Barrier(comm);
        }
#endif
    }
    par_CSRMatrixTranspose(* P_mat, * R_mat);
    {// 需要修改限制矩阵的值
        for (idx_t i = 0; i < R_mat->diag.nrows; i ++) {
            const idx_t num = R_mat->diag.row_ptr[i + 1] - R_mat->diag.row_ptr[i]
                            + R_mat->offd.row_ptr[i + 1] - R_mat->offd.row_ptr[i];
            const setup_t v = 1.0 / (setup_t) num;
            for (idx_t p = R_mat->diag.row_ptr[i]; p < R_mat->diag.row_ptr[i + 1]; p ++) {
                assert(R_mat->diag.vals[p] == 1.0);
                R_mat->diag.vals[p] = v;
            }
            for (idx_t p = R_mat->offd.row_ptr[i]; p < R_mat->offd.row_ptr[i + 1]; p ++) {
                assert(R_mat->offd.vals[p] == 1.0);
                R_mat->offd.vals[p] = v;
            }
        }
#ifdef DEBUG
        for (int p = 0; p < num_procs; p++) {
            if (my_pid == p) {
                printf("Proc %d Rstr diag %d %d %d offd %d %d %d\n", my_pid, R_mat->diag.nrows, R_mat->diag.ncols, R_mat->diag.nnz,
                    R_mat->offd.nrows, R_mat->offd.ncols, R_mat->offd.nnz);
                for (idx_t i = 0; i < R_mat->diag.nrows; i++) {
                    printf("  %d : ", R_mat->beg_row + i);
                    for (idx_t p = R_mat->diag.row_ptr[i]; p < R_mat->diag.row_ptr[i + 1]; p++)
                        printf("%d ", R_mat->beg_col + R_mat->diag.col_idx[p]);
                    printf(" | ");
                    for (idx_t p = R_mat->offd.row_ptr[i]; p < R_mat->offd.row_ptr[i + 1]; p++)
                        printf("%d ", R_mat->col_map_offd[R_mat->offd.col_idx[p]]);
                    printf("\n");
                }
            }
            MPI_Barrier(comm);
        }
#endif
    }
    delete [] glb_map_f2c;

#else // 最后的清扫只发生在进程内
    // 处理那些遗漏的未聚合点
    for (idx_t loc_fi = 0; loc_fi < my_num_F; loc_fi++) {
        if (loc_map_f2c[loc_fi] == UNAGGREGATED) {// 该点仍未被聚合
#if 1
            setup_t max_infl = 0.0;// 对我有最大影响的邻居的影响
            idx_t max_ngb = -1;// 对我有最大影响的邻居序号
            for (idx_t p = diag.row_ptr[loc_fi]; p < diag.row_ptr[loc_fi + 1]; p ++) {
                const idx_t loc_fj = diag.col_idx[p];
                if (loc_fj != loc_fi && // 该邻居不是自己
                    iso_wells.find(loc_fj + par_fine_mat->beg_row) == iso_wells.end() && // 该邻居不是孤立点
                    loc_map_f2c[loc_fj] != UNAGGREGATED) // 该邻居已被聚合
                {
                    setup_t infl = calc_blk_norm(diag.vals + p * e_size);// 计算该邻居对我的影响值
                    if (max_infl < infl) {
                        max_infl = infl;
                        max_ngb  = loc_fj;
                    }
                }
            }
            // 在有除了自身外的邻居的情况下，该点不可能所有邻居都未聚合！！否则它在之前就应该被发现和处理了！
            if (max_ngb == -1) {// 一个合适的聚合目标也没找到，则该点必然除了自身外没有其它邻居
                if (diag.row_ptr[loc_fi + 1] - diag.row_ptr[loc_fi] != 1) {
                    printf("Warning !! Proc %d glb_i %d num_ngbs %d max_ngb %d\n", my_pid, loc_fi + par_fine_mat->beg_row,
                        diag.row_ptr[loc_fi + 1] - diag.row_ptr[loc_fi], max_ngb);
                }
                // assert(diag.row_ptr[loc_fi + 1] - diag.row_ptr[loc_fi] == 1);
                // assert(diag.col_idx[diag.row_ptr[loc_fi]] == loc_fi);
            } else {// 找到合适目标后并入
                loc_map_f2c[loc_fi] = loc_map_f2c[max_ngb];// 归入该邻居
            }
#else
            idx_t unagg_ngb_cnt = 0;// 我的未聚合邻居数
            for (idx_t p = diag.row_ptr[loc_fi]; p < diag.row_ptr[loc_fi + 1]; p ++) {
                const idx_t loc_fj = diag.col_idx[p];
                if (loc_map_f2c[loc_fj] == UNAGGREGATED) // 该邻居未被聚合
                {
                    unagg_ngb_cnt ++;
                }
            }
            if (unagg_ngb_cnt > 1) {// 剩下的这些点可以聚合
                for (idx_t p = diag.row_ptr[loc_fi]; p < diag.row_ptr[loc_fi + 1]; p ++) {
                    const idx_t loc_fj = diag.col_idx[p];
                    if (loc_map_f2c[loc_fj] == UNAGGREGATED) // 该邻居未被聚合
                    {
                        loc_map_f2c[loc_fj] = my_num_C;
                    }
                }
                my_num_C ++;
            } else { assert(unagg_ngb_cnt == 1);// 我所有的邻居里就剩我没聚合了，找一个归宿
                setup_t max_infl = 0.0;
                idx_t max_ngb = loc_fi;
                for (idx_t p = diag.row_ptr[loc_fi]; p < diag.row_ptr[loc_fi + 1]; p ++) {
                    const idx_t loc_fj = diag.col_idx[p];
                    if (loc_fj != loc_fi && // 该邻居不是自己
                        iso_wells.find(loc_fj + par_fine_mat->beg_row) == iso_wells.end()) // 该邻居不是孤立点
                    {
                        setup_t infl = calc_blk_norm(diag.vals + p * e_size);// 计算该邻居对我的影响值
                        if (max_infl < infl) {
                            max_infl = infl;
                            max_ngb  = loc_fj;
                        }
                    }
                }
                if (max_infl == 0.0) {// 此时还尚未被聚合的，一定是只有一个邻居（即自己的）
                    assert(max_ngb == loc_fi);
                    // if (diag.row_ptr[loc_fi + 1] - diag.row_ptr[loc_fi] > 1) {
                    //     printf("Warning !! Proc %d glb_i %d ngb_cnt %d : %d (%.4e %d) %d (%.4e %d ) %d (%.4e %d)\n", my_pid, loc_fi + par_fine_mat->beg_row,
                    //         diag.row_ptr[loc_fi + 1] - diag.row_ptr[loc_fi],
                    //         diag.col_idx[diag.row_ptr[loc_fi]    ] + par_fine_mat->beg_row, calc_blk_norm(diag.vals + (diag.row_ptr[loc_fi]    )*e_size), loc_map_f2c[diag.col_idx[diag.row_ptr[loc_fi]    ]],
                    //         diag.col_idx[diag.row_ptr[loc_fi] + 1] + par_fine_mat->beg_row, calc_blk_norm(diag.vals + (diag.row_ptr[loc_fi] + 1)*e_size), loc_map_f2c[diag.col_idx[diag.row_ptr[loc_fi] + 1]],
                    //         diag.col_idx[diag.row_ptr[loc_fi] + 2] + par_fine_mat->beg_row, calc_blk_norm(diag.vals + (diag.row_ptr[loc_fi] + 2)*e_size), loc_map_f2c[diag.col_idx[diag.row_ptr[loc_fi] + 2]] );
                    // }
                    assert(diag.row_ptr[loc_fi + 1] - diag.row_ptr[loc_fi] == 1);
                } else {
                    loc_map_f2c[loc_fi] = loc_map_f2c[max_ngb];// 归入该邻居
                }
            }
#endif
        }
    }

    // 构造插值矩阵
    P_mat->diag.nrows = my_num_F; P_mat->offd.nrows = my_num_F;// 非对角部分行数照常
    P_mat->diag.ncols = my_num_C; P_mat->offd.ncols = 0;// 非对角部分列数为0
    P_mat->diag.row_ptr = new idx_t [my_num_F + 1]; P_mat->diag.row_ptr[0] = 0;
    P_mat->offd.row_ptr = new idx_t [my_num_F + 1]; P_mat->offd.row_ptr[0] = 0;

    idx_t twspace [omp_get_max_threads()];
    #pragma omp parallel
    {
        int tid = omp_get_thread_num();
        int  nt = omp_get_num_threads();
        idx_t tbeg = -1, tend = -1;
        thread_load_balance(my_num_F, nt, tid, tbeg, tend);

        idx_t tcnt = 0;
        for (idx_t loc_fi = tbeg; loc_fi < tend; loc_fi ++) {
            idx_t loc_ci = loc_map_f2c[loc_fi];
            if (loc_ci == UNAGGREGATED) {
                P_mat->diag.row_ptr[loc_fi + 1] = tcnt;// 插值矩阵里该行为空
            } else {
                tcnt ++;
                P_mat->diag.row_ptr[loc_fi + 1] = tcnt;// plain
            }
            P_mat->offd.row_ptr[loc_fi + 1] = 0;
        }
        twspace[tid] = tcnt;

        #pragma omp barrier
        if (tid > 0) {
            idx_t prev_t_sum = 0;// prefix sum
            for (int prev_t = 0; prev_t < tid; prev_t ++) prev_t_sum += twspace[prev_t];
            for (idx_t loc_fi = tbeg; loc_fi < tend; loc_fi ++)
                P_mat->diag.row_ptr[loc_fi + 1] += prev_t_sum;
        } else {// tid == 0
            idx_t tot_t_sum = 0;// all threads sum
            for (int t = 0; t < nt; t ++) tot_t_sum += twspace[t];
            P_mat->diag.nnz = tot_t_sum;
            P_mat->diag.col_idx = new idx_t  [P_mat->diag.nnz];
            P_mat->diag.vals    = new setup_t[P_mat->diag.nnz];
        }
        
        #pragma omp barrier
        
        for (idx_t loc_fi = tbeg; loc_fi < tend; loc_fi ++) {
            idx_t loc_ci = loc_map_f2c[loc_fi];
            if (loc_ci != UNAGGREGATED) {
                const idx_t off = P_mat->diag.row_ptr[loc_fi];
                P_mat->diag.col_idx[off] = loc_ci;
                P_mat->diag.vals   [off] = 1.0;// plain
            }
        }
    }
    // 转置构造限制矩阵
    CSRMatrixTranspose(P_mat->diag, R_mat->diag);
    R_mat->offd.nrows = R_mat->diag.nrows;// 非对角部分行数照常
    R_mat->offd.ncols = 0;// 非对角部分列数为0
    R_mat->offd.row_ptr = new idx_t [R_mat->offd.nrows + 1]; R_mat->offd.row_ptr[0] = 0;
    // 需要修改限制矩阵的值
    #pragma omp parallel for schedule(static)
    for (idx_t i = 0; i < R_mat->diag.nrows; i++) {
        const idx_t num = R_mat->diag.row_ptr[i + 1] - R_mat->diag.row_ptr[i]; assert(num > 0);
        const setup_t v = 1.0 / (setup_t) num;
        for (idx_t p = R_mat->diag.row_ptr[i]; p < R_mat->diag.row_ptr[i + 1]; p++) {
            assert(R_mat->diag.vals[p] == 1.0);
            R_mat->diag.vals[p] = v;
        }
        R_mat->offd.row_ptr[i + 1] = 0;
    }
#endif // 最后的清扫是否可跨进程
#ifdef DEBUG
        for (int p = 0; p < num_procs; p++) {
            if (my_pid == p) {
                printf("Proc %d Rstr diag %d %d %d offd %d %d %d\n", my_pid, R_mat->diag.nrows, R_mat->diag.ncols, R_mat->diag.nnz,
                    R_mat->offd.nrows, R_mat->offd.ncols, R_mat->offd.nnz);
                for (idx_t i = 0; i < R_mat->diag.nrows; i++) {
                    printf("  %d : ", R_mat->beg_row + i);
                    for (idx_t p = R_mat->diag.row_ptr[i]; p < R_mat->diag.row_ptr[i + 1]; p++)
                        printf("%d ", R_mat->beg_col + R_mat->diag.col_idx[p]);
                    printf(" | ");
                    for (idx_t p = R_mat->offd.row_ptr[i]; p < R_mat->offd.row_ptr[i + 1]; p++)
                        printf("%d ", R_mat->col_map_offd[R_mat->offd.col_idx[p]]);
                    printf("\n");
                }
            }
            MPI_Barrier(comm);
        }
        // MPI_Barrier(comm);
        // MPI_Abort(comm, -1212);
#endif
// #ifdef DEBUG
    const idx_t my_num_I = my_num_F - P_mat->diag.nnz - P_mat->offd.nnz;// 独立点个数
    idx_t tot_num_I = 0;
    MPI_Reduce(& my_num_I, & tot_num_I, 1, mpi_idx_type, MPI_SUM, 0, comm);
    if (my_pid == 0) printf("Glb #I %d\n", tot_num_I);
// #endif
    delete [] loc_map_f2c;
    delete [] been_in_Q;
    return C_rows_partits;
}

template int* naive_neighbor(const par_CSRMatrix_I32F64F64_DOF1 *, par_CSRMatrix_I32F64F64_DOF1 * &, par_CSRMatrix_I32F64F64_DOF1 * &);
template int* naive_neighbor(const par_CSRMatrix_I32F64F64_DOF2 *, par_CSRMatrix_I32F64F64_DOF1 * &, par_CSRMatrix_I32F64F64_DOF1 * &);
template int* naive_neighbor(const par_CSRMatrix_I32F64F64_DOF3 *, par_CSRMatrix_I32F64F64_DOF1 * &, par_CSRMatrix_I32F64F64_DOF1 * &);
template int* naive_neighbor(const par_CSRMatrix_I32F64F64_DOF4 *, par_CSRMatrix_I32F64F64_DOF1 * &, par_CSRMatrix_I32F64F64_DOF1 * &);

template int* naive_neighbor(const par_CSRMatrix_I32F32F32_DOF1 *, par_CSRMatrix_I32F32F32_DOF1 * &, par_CSRMatrix_I32F32F32_DOF1 * &);
template int* naive_neighbor(const par_CSRMatrix_I32F32F32_DOF2 *, par_CSRMatrix_I32F32F32_DOF1 * &, par_CSRMatrix_I32F32F32_DOF1 * &);
template int* naive_neighbor(const par_CSRMatrix_I32F32F32_DOF3 *, par_CSRMatrix_I32F32F32_DOF1 * &, par_CSRMatrix_I32F32F32_DOF1 * &);
template int* naive_neighbor(const par_CSRMatrix_I32F32F32_DOF4 *, par_CSRMatrix_I32F32F32_DOF1 * &, par_CSRMatrix_I32F32F32_DOF1 * &);

